import yfinance as yf
import pandas as pd
import matplotlib.pyplot as plt

# 获取股票数据
def get_stock_data(ticker, start_date, end_date):
    data = yf.download(ticker, start=start_date, end=end_date)
    return data

# 计算简单移动平均线（SMA）
def calculate_sma(data, window):
    data['SMA'] = data['Close'].rolling(window=window).mean()
    return data

# 计算指数移动平均线（EMA）
def calculate_ema(data, span):
    data['EMA'] = data['Close'].ewm(span=span, adjust=False).mean()
    return data

# 绘制股票价格和移动平均线图表
def plot_chart(data):
    plt.figure(figsize=(12, 6))
    plt.plot(data['Close'], label='Close Price', color='blue')
    plt.plot(data['SMA'], label='SMA', color='red')
    plt.plot(data['EMA'], label='EMA', color='green')
    plt.title('Stock Price and Moving Averages')
    plt.xlabel('Date')
    plt.ylabel('Price')
    plt.legend()
    plt.show()

# 主函数
def main():
    ticker = '588080.SS'  # 股票代码（苹果公司）
    start_date = '2024-01-01'  # 开始日期
    end_date = '2024-05-20'  # 结束日期
    window = 20  # SMA 窗口大小
    span = 20  # EMA 平滑系数

    # 获取股票数据
    stock_data = get_stock_data(ticker, start_date, end_date)

    # 计算移动平均线
    stock_data = calculate_sma(stock_data, window)
    stock_data = calculate_ema(stock_data, span)

    # 绘制图表
    plot_chart(stock_data)

if __name__ == "__main__":
    main()
